import collections
# 子数组求和问题
def numSubarraysWithSum(A, S):
    indexes = [-1] + [ix for ix, v in enumerate(A) if v] + [len(A)] # 这个两头加的太巧妙了
    print(indexes)
    ans = 0
    if S == 0:  # 目标是0，那我就看连续的零的个数有多少个
        for x in range(len(indexes) - 1):
            w = indexes[x + 1] - indexes[x] - 1 # 表示原数组中连续两个1得间隔
            ans = ans + w*(w + 1)/2
        return ans

    for i in range(1, len(indexes) - S):
        j = i + S - 1   # 中间应该有S个1，找到下一个1的位置
        left = indexes[i] - indexes[i - 1]
        right = indexes[j + 1] - indexes[j]
        ans = ans + left * right
    return ans

# def numSubarraysWithSum(A, S):
#         P = [0]
#         for x in A: P.append(P[-1] + x)
#         count = collections.Counter()
#
#         ans = 0
#         for x in P:
#             ans += count[x]
#             count[x + S] += 1
#
#         return ans

if __name__ == '__main__':
    numSubarraysWithSum([1,0,1,0,1],2)